package zju.dgm.service.loadbalance;

import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.recipes.cache.PathChildrenCache;
import org.apache.curator.framework.recipes.cache.PathChildrenCacheEvent;
import org.apache.curator.framework.recipes.cache.PathChildrenCacheListener;
import org.apache.curator.framework.recipes.locks.InterProcessReadWriteLock;
import org.apache.curator.utils.ZKPaths;
import org.apache.zookeeper.CreateMode;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import zju.dgm.dto.QueryStat;
import zju.dgm.dto.RegionEntry;
import zju.dgm.dto.ServerInfo;
import zju.dgm.service.election.ElectionService;
import zju.dgm.service.election.MasterUtilProvider;
import zju.dgm.service.name.NameService;
import zju.dgm.service.ServiceEntry;
import zju.dgm.service.region.RegionUtilProvider;
import zju.dgm.util.CacheManager;
import zju.dgm.util.Serializer;
import zju.dgm.util.SyncUtil;


import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;


/**
 * @author lzy
 */
@Component
public class LoadBalanceService extends ServiceEntry implements LoadBalanceUtilProvider {

    private final static Logger logger = LoggerFactory.getLogger(LoadBalanceService.class);

    public static final CountDownLatch finLatch = new CountDownLatch(1);


    final long LB_UPDATE_TIME = 120;
    final double UPDATE_RATIO = 0.75;

    @Autowired
    LoadBalancePath path;

    @Autowired
    RegionUtilProvider regionUtilProvider;

    @Autowired
    MasterUtilProvider masterUtilProvider;


    @Autowired
    private CuratorFramework client;

    @Autowired
    private CacheManager cacheManager;


    QueryStat stat = new QueryStat();


    @Override
    public String getRootPath() {
        return path.getRootPath();
    }

    @Override
    public boolean needRootPath() {
        return true;
    }

    @Override
    protected CountDownLatch[] getPrerequisites() {
        return new CountDownLatch[]{
                NameService.finLatch,
                ElectionService.finLatch,
//                RegionServer.finLatch
        };
    }


    @Override
    protected void registerListener() {


    }

    private void realBuildInfoListener() {
        PathChildrenCache lbInfoCache = cacheManager.getPathChildrenCache(path.getInfoPath());
        lbInfoCache.getListenable().addListener(new PathChildrenCacheListener() {
            @Override
            public void childEvent(CuratorFramework curatorFramework, PathChildrenCacheEvent event) throws Exception {
                if (event.getType() == PathChildrenCacheEvent.Type.CHILD_REMOVED) {
                    logger.info("节点{}断连", event.getData().getPath());
                }
            }
        });
    }

    @Override
    protected void start() {


        try {
            // 重置信息，清空残留
            client.setData().forPath(path.getRootPath(), null);
            ZKPaths.mkdirs(client.getZookeeperClient().getZooKeeper(), path.getInfoPath());

            logger.info("正在启动负载均衡服务");
            String selfInfoPath = ZKPaths.makePath(path.getInfoPath(), SyncUtil.serverName);
            realBuildInfoListener();
            // 创建自己的信息路径,必须要是临时的
            client.create().withMode(CreateMode.EPHEMERAL).forPath(selfInfoPath, null);


            // master工作线程
            masterUtilProvider.addMasterTask(() -> {
                logger.info("master节点正在维护负载平衡信息");
                // master节点维护的负载平衡信息
                HashMap<String, HashMap<ServerInfo, Integer>> masterStatMap = null;

                InterProcessReadWriteLock rwLock = new InterProcessReadWriteLock(client, path.getRwLockPath());
                try {


                    while (true) {
                        List<String> servers = client.getChildren().forPath(path.getInfoPath());
                        // 用于暂存每一轮的负载信息
                        HashMap<String, HashMap<ServerInfo, Integer>> tempMap = new HashMap<>();

                        Map<String, List<ServerInfo>> map = regionUtilProvider.getRegionMap();
                        if (map != null && map.size() != 0) {
                            // 强制获取分布表同步并初始化一次region信息
                            map.forEach((k, v) -> {
                                HashMap<ServerInfo, Integer> cntMap = new HashMap<>();
                                for (ServerInfo serverInfo : v) {
                                    cntMap.put(serverInfo, 0);
                                }
                                tempMap.put(k, cntMap);
                            });
                        }


                        logger.info("在master，经过初始化后的tempMap的大小：{}", tempMap.size());


                        for (String serverId : servers) {
                            byte[] stat = client.getData().forPath(ZKPaths.makePath(path.getInfoPath(), serverId));
                            if (stat == null || stat.length == 0) {
                                logger.info("region server {} 尚未写入", serverId);
                                continue;
                            }// 改server还没有写入

                            // 可能因为没有表，因此结果是[]两个字节
                            ConcurrentHashMap<String, PriorityQueue<RegionEntry>> regionMap = Serializer.getQueueFromDTO(
                                    Serializer.getObjectFromJSON(
                                            new String(stat),
                                            List.class
                                    )
                            );
                            if (regionMap == null) {
                                continue;
                            }
                            logger.info("{} 的map:{}", serverId, regionMap);

                            // 针对每一个region，累加数值
                            for (Map.Entry<String, PriorityQueue<RegionEntry>> entry : regionMap.entrySet()) {
                                if (!tempMap.containsKey(entry.getKey())) {
                                    continue;
                                }
                                HashMap<ServerInfo, Integer> cntMap = tempMap.get(entry.getKey());
                                for (RegionEntry regionEntry : entry.getValue()) {
                                    if (!cntMap.containsKey(regionEntry.getServer())) {
                                        continue;
                                    }

                                    // 核心部分
                                    int prevVal;
                                    if (masterStatMap == null
                                            || masterStatMap.get(entry.getKey()) == null
                                            || masterStatMap.get(entry.getKey()).get(regionEntry.getServer()) == null) {
                                        prevVal = 0;
                                    } else {
                                        prevVal = masterStatMap.get(entry.getKey()).get(regionEntry.getServer());
                                    }
                                    cntMap.put(
                                            regionEntry.getServer(),
                                            cntMap.get(regionEntry.getServer()) + regionEntry.getQTimes() - prevVal
                                    );
                                }
                            }
                        }

                        // 更新时用写锁获取同步状态
                        rwLock.writeLock().acquire();
                        byte[] toBeWritten = Serializer.getJSONString(
                                Serializer.getDTOFromMap(
                                        tempMap
                                )
                        ).getBytes(StandardCharsets.UTF_8);
                        logger.info("master 写出了一次经维护后的负载信息，其字节长度为：{}", toBeWritten.length);
                        client.setData().forPath(path.getRootPath(), toBeWritten);

                        rwLock.writeLock().release();

                        // 更新线程持有的长期状态
                        masterStatMap = tempMap;
                        Thread.sleep((long) (LB_UPDATE_TIME * 1000 * UPDATE_RATIO));

                    }
                } catch (Exception e) {
                    e.printStackTrace();
                    logger.info("master节点维护负载平衡INFO时错误");
                } finally {
                    try {
                        if (rwLock.writeLock().isAcquiredInThisProcess()) {
                            rwLock.writeLock().release();
                        }
                    } catch (Exception e) {
                        e.printStackTrace();
                    }
                }

            });

//             region工作线程
            SyncUtil.tpe.submit(() -> {
                InterProcessReadWriteLock rwLock = new InterProcessReadWriteLock(client, path.getRwLockPath());
                logger.info("当前节点正在维护节点信息");
                while (true) {
                    try {
                        // 1 获取读锁，从master节点读信息来进行初始化
                        rwLock.readLock().acquire();
                        byte[] masterStat = client.getData().forPath(path.getRootPath());
                        rwLock.readLock().release();
                        if (masterStat == null || masterStat.length == 0) {
                            stat.defaultInit(regionUtilProvider.getRegionMap());// 用表集群信息默认初始化
                        } else {// 用master维护的信息初始化
                            ConcurrentHashMap<String, PriorityQueue<RegionEntry>> rebuildMap =
                                    Serializer.getQueueFromDTO(
                                            Serializer.getObjectFromJSON(
                                                    new String(masterStat),
                                                    List.class
                                            )
                                    );
                            logger.info("读取到了master维护的负载信息:{}", rebuildMap);
                            if (rebuildMap == null) {
                                stat.defaultInit(regionUtilProvider.getRegionMap());
                            } else {
                                // 读master来更新
                                stat.updateStats(rebuildMap);
                            }

                        }

                        Thread.sleep(LB_UPDATE_TIME * 1000);
                        // 2 向自己的节点写入自身负载信息
                        String toBeWritten = Serializer.getJSONString(
                                Serializer.getDTOFromQueue(stat.getTableStats())
                        );
                        logger.info("{}写出了一次负载信息：{}", SyncUtil.serverName, toBeWritten);
                        client.setData().forPath(selfInfoPath, toBeWritten.getBytes(StandardCharsets.UTF_8));


                    } catch (Exception e) {
                        e.printStackTrace();
                        logger.error("服务节点负载平衡线程错误");
                        Thread.sleep(LB_UPDATE_TIME * 1000);
                    }

                    // 3 周期性等待，将其放在最后能确保负载平衡服务运行初期就能获取并记录负载信息

                }
            });
            finLatch.countDown();// lb服务加载完成
            logger.info("负载均衡服务启动完成");


        } catch (Exception e) {
            logger.error("启动负载均衡服务过程中出错");
            e.printStackTrace();
        }

    }

    @Override
    public ServerInfo getRequestingServer(String tableName) throws Exception {
//        Map<String, List<ServerInfo>> info = regionUtilProvider.getRegionMap();
//        if (!info.containsKey(tableName) || info.get(tableName).size() == 0) {
//            return null;
//        } else {
//            return info.get(tableName).get(0);
//        }
        return tryGetRequestingServer(tableName, 3);
    }

    @Override
    public ServerInfo tryGetRequestingServer(String tableName, int times) throws Exception {

        // 循环一定次数
        for (int i = 0; i < times - 1; i++) {
            try {
                return stat.getTablePos(tableName);
            } catch (Exception error) {
                error.printStackTrace();
                // 如果出现了AssertionError,说明server状况已被更新
                stat.syncInfoIncrement(regionUtilProvider.getRegionMap());
            }
        }
        ServerInfo rst = stat.getTablePos(tableName);
        if (rst == null) {
            Map<String, List<ServerInfo>> info = regionUtilProvider.getRegionMap();
            if (!info.containsKey(tableName) || info.get(tableName).size() == 0) {
                return null;
            } else {
                return info.get(tableName).get(0);
            }
        }
        return rst;
    }

    @Override
    public void touch(String table) {
        stat.touch(table);
    }


}
